package com.lovely602.langchain4j.embedding.controller;


import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.output.Response;
import dev.langchain4j.store.embedding.EmbeddingSearchRequest;
import dev.langchain4j.store.embedding.EmbeddingSearchResult;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.filter.MetadataFilterBuilder;
import io.qdrant.client.QdrantClient;
import io.qdrant.client.grpc.Collections;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;

/**
 * 知识出处，https://docs.langchain4j.dev/tutorials/rag#embedding-store
 *
 * @author lizhixing
 */
@Slf4j
@RestController
@RequestMapping("/embedding")
public class EmbeddingController {

    @Autowired
    private EmbeddingModel embeddingModel;

    @Autowired
    private QdrantClient qdrantClient;

    @Autowired
    private EmbeddingStore<TextSegment> embeddingStore;

    /**
     * 文本向量化测试，看看形成向量后的文本
     * http://localhost:9012/embedding/embed
     * @return
     */
    @GetMapping(value = "/embed")
    public String embed() {
        String prompt = """
                 他乡剪烛
                灯火随风转，
                吴音入耳频。
                惊看犹是旧时人。
                一揖忽收千里雨，
                拍肩同拂十年尘。
                未添华发泪先新。
                剪烛话长夜，
                举杯邀早春。
                他乡今夕作比邻。
                """;
        Response<Embedding> embeddingResponse = embeddingModel.embed(prompt);

        System.out.println(embeddingResponse);

        return embeddingResponse.content().toString();
    }

    /**
     * 新建向量数据库实例和创建索引：test-qdrant
     * 类似mysql create database test-qdrant
     * http://localhost:8899/embedding/createCollection
     */
    @GetMapping(value = "/createCollection")
    public void createCollection() {
        var vectorParams = Collections.VectorParams.newBuilder()
                .setDistance(Collections.Distance.Cosine)
                .setSize(1024)
                .build();
        qdrantClient.createCollectionAsync("test-qdrant", vectorParams);
    }

    /**
     * 往向量数据库新增文本记录
     */
    @GetMapping(value = "/add")
    public String add() {
        String prompt = """
                 他乡剪烛
                灯火随风转，
                吴音入耳频。
                惊看犹是旧时人。
                一揖忽收千里雨，
                拍肩同拂十年尘。
                未添华发泪先新。
                剪烛话长夜，
                举杯邀早春。
                他乡今夕作比邻。
                """;
        TextSegment segment1 = TextSegment.from(prompt);
        segment1.metadata().put("author", "lzx");
        Embedding embedding1 = embeddingModel.embed(segment1).content();
        String result = embeddingStore.add(embedding1, segment1);

        System.out.println(result);

        return result;
    }

    /**
     * 搜索向量数据库
     * http://localhost:8899/embedding/query1
     */
    @GetMapping(value = "/query1")
    public String query1() {
        Embedding queryEmbedding = embeddingModel.embed("他乡剪烛说的是什么").content();
        EmbeddingSearchRequest embeddingSearchRequest = EmbeddingSearchRequest.builder()
                .queryEmbedding(queryEmbedding)
                .maxResults(1)
                .build();
        EmbeddingSearchResult<TextSegment> searchResult = embeddingStore.search(embeddingSearchRequest);
        System.out.println(searchResult.matches().get(0).embedded().text());
        return searchResult.matches().get(0).embedded().text();
    }

    /**
     * 搜索向量数据库，只搜索作者为lzx的
     */
    @GetMapping(value = "/query2")
    public String query2() {
        Embedding queryEmbedding = embeddingModel.embed("他乡剪烛").content();

        EmbeddingSearchRequest embeddingSearchRequest = EmbeddingSearchRequest.builder()
                .queryEmbedding(queryEmbedding)
                .filter(MetadataFilterBuilder.metadataKey("author").isEqualTo("lzx"))
                .maxResults(1)
                .build();

        EmbeddingSearchResult<TextSegment> searchResult = embeddingStore.search(embeddingSearchRequest);

        System.out.println(searchResult.matches().get(0).embedded().text());
        return searchResult.matches().get(0).embedded().text();
    }

}
